library(librarian)
shelf(ggpubr,
      patchwork,
      ggrepel,
      tidyverse)

theme_jpub <- theme_classic(base_size = 6, base_line_size = 1) +
  theme(plot.margin = margin(), legend.box.margin = margin(),
        plot.title.position = 'plot')

prepare_volc <- function(df, group1 = 'group1', group2 = 'group2', fc_thres = 1, pval_thres = .05,
                       highlights = NULL, force = FALSE){
  if(is.null(df$gene)){df <- as_tibble(df, rownames = 'gene')}
  
  #group1 <- str_c('Up in ',group1)
  #group2 <- str_c('Up in ',group2)
  
  df <- df |>
    mutate(type = case_when(avg_log2FC > fc_thres & p_val_adj < .05 ~ group1,
                            avg_log2FC < -fc_thres & p_val_adj < .05 ~ group2,
                            .default = 'NS') |> fct_relevel(group2, 'NS'))
  
  symm_x_lim <- ceiling(df$avg_log2FC |> abs() |> max())
  
  df |>
    dplyr::count(type) |> print()
  
  if(!is.null(highlights)){
    seek_name <- df |>
      filter(gene %in% highlights)
    if (!force) {
      seek_name <- seek_name |>
        filter(type != 'NS')
    }
  } else {seek_name <- tibble()}
  
  if(nrow(seek_name) == 0){
    warning('No interested genes upregulated or specified. Will highlight top-DEGs.')
    top.fc <- df |>
      filter(type != 'NS') |>
      group_by(type) |>
      mutate(abs_fc = abs(avg_log2FC)) |>
      slice_max(abs_fc, n = 5, with_ties = F) |>
      pull(gene)
    
    top.sig <- df |>
      filter(type != 'NS') |>
      group_by(type) |>
      mutate(abs_fc = abs(avg_log2FC)) |>
      slice_min(p_val_adj, n = 5, with_ties = F) |>
      pull(gene)
    
    seek_name <- df |>
      filter(gene %in% c(top.sig, top.fc))
  }
  
  list(df, seek_name, symm_x_lim)
}

#' Plot volcano for generic & slides 
#'
#' @param df data.frame from Seurat::FindMarker() output.
#' @param group1 character vector of ident.1 name, for legend.
#' @param fc_thres log2fc threshold for dash line threshold, 1 by default.
#' @param pval_thres adjusted p value threshold for dash line threshold, 0.05 by default.
#' @param highlights optional character string of gene names wanted to highlight. Will auto highlight genes if leave empty.
#' @param force Boolean, highlight `highlights` genes anyway, even they are not significantly upregulated? FALSE by default.
#'
#' @return A ggolot object.
#' @export
#'
#' @examples 
#' pbmc_small |>
#'   FindMarker(ident.1 = 0) |>
#'   plot_bill_volc('clus0')
#'   
plot_bill_volc <- function(df, group1, group2, fc_thres = 1, pval_thres = .05,
                          highlights = NULL, force = FALSE){
  df_list <- prepare_volc(df, group1, group2, fc_thres, pval_thres,
               highlights, force)
  seek_name <- df_list[[2]]
  symm_x_lim <- df_list[[3]]
  
  df_list[[1]] |>
    ggplot(aes(avg_log2FC, -log10(p_val_adj), color = type)) +
    geom_point(size = .5, alpha = .3) +
    geom_vline(xintercept = c(-1,1), linetype = 'dashed') +
    geom_hline(yintercept = 1.3, linetype = 'dashed') +
    geom_point(data = seek_name, color = 'purple') +
    geom_text_repel(data = seek_name, aes(label = gene), color = 'purple') +
    scale_color_manual(values = c('blue','grey','red') ) +
    theme_pubr() +
    theme(legend.position = 'top', plot.margin = margin(), legend.box.margin = margin()) +
    expand_limits(x = c(-symm_x_lim,symm_x_lim))
}

# TODO: refactor two volc func to avoid replicate self! =========

#' Plot volcano for journal-publication 
#'
#' @param df data.frame from Seurat::FindMarker() output.
#' @param group1 character vector of ident.1 name, for legend.
#' @param fc_thres log2fc threshold for dash line threshold, 1 by default.
#' @param pval_thres adjusted p value threshold for dash line threshold, 0.05 by default.
#' @param highlights optional character string of gene names wanted to highlight. Will auto highlight genes if leave empty.
#' @param force Boolean, highlight `highlights` genes anyway, even they are not significantly upregulated? FALSE by default.
#'
#' @return A ggolot object.
#' @export
#'
#' @examples 
#' pbmc_small |>
#'   FindMarker(ident.1 = 0) |>
#'   plot_pub_volc('clus0')
#'   
plot_pub_volc <- function(df, group1, group2, fc_thres = 1, pval_thres = .05,
                           highlights = NULL, force = FALSE){
  df_list <- prepare_volc(df, group1, group2, fc_thres, pval_thres,
                          highlights, force)
  seek_name <- df_list[[2]]
  symm_x_lim <- df_list[[3]]
  
  df_list[[1]] |>
    ggplot(aes(avg_log2FC, -log10(p_val_adj), color = type)) +
    geom_point(size = .3, alpha = .3) +
    geom_vline(xintercept = c(-1,1), linetype = 'dashed') +
    geom_hline(yintercept = 1.3, linetype = 'dashed') +
    geom_point(data = seek_name, size = 1, color = 'purple') +
    geom_text_repel(data = seek_name, aes(label = gene), color = 'purple', size = 2) +
    scale_color_manual(values = c('blue','grey','red') ) +
    theme_classic(base_size = 6, base_line_size = 1) +
    theme(legend.position = 'top', plot.margin = margin(), legend.box.margin = margin()) +
    expand_limits(x = c(-symm_x_lim,symm_x_lim))
}

plot_enrichment <- function(df, base_col = 'red', n = 10, descr_wrap = 40){
  df |>
    filter(p.adjust < .05) |>
    slice_min(p.adjust, n = n, with_ties = F) |>
    mutate(Description = str_wrap(Description, descr_wrap) |>
             fct_reorder(Count)) |>
    ggplot(aes(Description, Count, fill = p.adjust)) +
    geom_col() +
    coord_flip() +
    theme_pubr(legend = 'right') +
    labs_pubr() +
    scale_fill_gradient(low = base_col, high = 'black') +
    theme(plot.title.position = 'plot') +
    ylab('Gene count')
}

publish_enrichment <- function(df, base_col = 'red', n = 10){
  df |>
    filter(p.adjust < .05) |>
    slice_min(p.adjust, n = n, with_ties = F) |>
    mutate(Description = str_wrap(Description, 30) |>
             fct_reorder(Count)) |>
    ggplot(aes(Description, Count, fill = p.adjust)) +
    geom_col() +
    coord_flip() +
    theme_classic(base_size = 6, base_line_size = 1) +
    scale_fill_gradient(low = base_col, high = 'black') +
    theme(plot.title.position = 'plot') +
    ylab('Gene count')
}

#' Calculate fraction confidence interval of subtype in group
#' 
#' Useful for ggplot col plot with error bars.
#'
#' @param df A tibble with `group` and `subtype` columns, typically a scRNA-seq meta data.
#' @param group colname of column containing variable of interest (e.g. WT/KO), need to have more than 1 types.
#' @param subtype colname of column containing variable of cell types (e.g. CD4T/CD8T/NK etc.)
#'
#' @return A tibble added some columns: 
#' - `fraction`: subtype fraction in group
#' - `conf.low` & `conf.high`: 95% CI of fraction, by `t.test`
#' - `n0`: sum of group count - subtype count
#' @export
#'
#' @examples 
#' meta_kera |>
#'   calc_frac_conf_on_grouped_count(orig.ident, seurat_clusters)
#' 
calc_frac_conf_on_grouped_count <- function(df, group, subtype){
  df |>
    dplyr::count({{ group }}, {{ subtype }}) |>
    group_by({{ group }}) |>
    mutate(n0 = sum(n) - n,
           fraction = n / sum(n)) |>
    rowwise() |>
    mutate(conf.low = t.test(c(rep(1,n),rep(0,n0)))$conf.int[1],
           conf.high = t.test(c(rep(1,n),rep(0,n0)))$conf.int[2]) |>
    ungroup()
}


#' Test significance for fraction of subtype in group
#'
#' @param df A tibble with `group` and `subtype` columns, typically a scRNA-seq meta data.
#' @param group colname of column containing variable of interest (e.g. WT/KO), need to have 2 and only 2 types.
#' @param subtype colname of column containing variable of cell types (e.g. CD4T/CD8T/NK etc.)
#'
#' @return A tibble with 2 columns: 
#' - `subtype`: defined `subtype`
#' - `p.value`: p value calculated from t.test
#' @export
#'
#' @examples
#' meta_kera |>
#'   test_on_grouped_count(orig.ident, seurat_clusters)
#'   
test_on_grouped_count <- function(df, group, subtype){
  df |>
    dplyr::count({{ group }}, {{ subtype }}) |>
    group_by({{ group }}) |>
    mutate(n0 = sum(n) - n) |>
    ungroup() |>
    pivot_wider(names_from = {{ group }}, values_from = c(n, n0)) |>
    set_names(c('subtype', 'n1_x', 'n1_y', 'n0_x', 'n0_y')) |>
    rowwise() |>
    mutate(
      seq1 = list(c(rep(1,n1_x),rep(0,n0_x))),
      seq2 = list(c(rep(1,n1_y),rep(0,n0_y))),
      p.value = t.test(seq1, seq2)$p.value) |>
    select(subtype, p.value) |>
    ungroup()
}

#' Compare cell type fraction changes between 2 conditions
#'
#' @param df A tibble with `group` and `subtype` columns, typically a scRNA-seq meta data.
#' @param group colname of column containing variable of interest (e.g. WT/KO), need to have 2 and only 2 types.
#' @param subtype colname of column containing variable of cell types (e.g. CD4T/CD8T/NK etc.)
#' @param var.1 identity class in `group` as numerator to calculate log2FC.
#' @param var.2 identity class in `group` as denominator for comparison.
#'
#' @return A tibble with 6 columns: 
#' - `subtype`: defined `subtype`
#' - `p.value`: p value calculated from t.test
#' - 2 columns named the same as `var.1` and `var.2`: their fraction value
#' - `log2fc_frac`: log2 of fold change of `var.1`/`var.2`
#' - `type`: 'NS''Decreased''Increased' for adjusted p value in `var.1`/`var.2`
#' @export
#'
#' @examples
#' meta_kera |>
#'   discov_frac_change(orig.ident, seurat_clusters)
#'   
discov_frac_change <- function(meta, group, subtype, var.1, var.2) {
  fmeta <- meta |>
    mutate(fvar = {{ group }}, ftype = {{ subtype }}, .keep = 'none')
  
  logfc_x <- fmeta |>
    calc_frac_conf_on_grouped_count(fvar, ftype) |>
    pivot_wider(names_from = fvar, values_from = fraction,
                id_cols = ftype) |>
    mutate(log2fc_frac = log2({{ var.1 }}/{{ var.2 }}))
  
  fmeta |>
    test_on_grouped_count(fvar, ftype) |>
    left_join(logfc_x, join_by(subtype == ftype)) |>
    mutate(p.value = p.adjust(p.value, method = 'bonferroni'),
           type = case_when(
             p.value < .05 & log2fc_frac > 0 ~ 'Increased',
             p.value < .05 & log2fc_frac < 0 ~ 'Decreased',
             .default = 'NS'),
           subtype = fct_reorder(subtype, log2fc_frac))
}


publish_pdf <- function(file, width = 50, height = 50){
  ggsave(file, width = width, height = height, units = 'mm')
}

#' Make dataframe `x` for ggpubr::stat_pvalue_manual(data = x)
#'
#' @param df The df containing all following variables, for adding layer of stat_pvalue_manual().
#' @param y The variable name of y in ggplot(aes()).
#' @param facets Required. The facet variable in main ggplot .
#' @param p.value The variable containing p value.
#' @param group.1 Character. The group 1 in testing.
#' @param group.2 Character. The group 2 in testing.
#'
#' @return A df for ggpubr::stat_pvalue_manual(data = x)
#' @export
#'
#' @examples
customize_pvalue <- function(df, y, facets, p.value, group.1, group.2){
  df |>
    group_by({{ facets }}) |>
    reframe(y.position = max({{ y }}) * 1.05,
            p = mean({{ p.value }}) |> signif(3),
            group1 = group.1,
            group2 = group.2)
}

add_manual_pvalue <- function(x){
  stat_pvalue_manual(data = x,label = 'p',hide.ns = 'p')}
